package com.hyacinth.cloudnote.modules.common.isolation;

import com.hyacinth.cloudnote.common.exception.AppRuntimeException;
import com.hyacinth.cloudnote.common.security.CurrentUserContextHolder;
import com.hyacinth.cloudnote.common.security.UserPrincipal;
import lombok.Data;
import lombok.experimental.Accessors;
import lombok.extern.slf4j.Slf4j;
import net.sf.jsqlparser.expression.BinaryExpression;
import net.sf.jsqlparser.expression.Expression;
import net.sf.jsqlparser.expression.HexValue;
import net.sf.jsqlparser.expression.Parenthesis;
import net.sf.jsqlparser.expression.operators.conditional.AndExpression;
import net.sf.jsqlparser.expression.operators.conditional.OrExpression;
import net.sf.jsqlparser.expression.operators.relational.*;
import net.sf.jsqlparser.parser.CCJSqlParserUtil;
import net.sf.jsqlparser.schema.Column;
import net.sf.jsqlparser.schema.Table;
import net.sf.jsqlparser.statement.Statement;
import net.sf.jsqlparser.statement.Statements;
import net.sf.jsqlparser.statement.delete.Delete;
import net.sf.jsqlparser.statement.insert.Insert;
import net.sf.jsqlparser.statement.select.*;
import net.sf.jsqlparser.statement.update.Update;
import org.hibernate.resource.jdbc.spi.StatementInspector;
import org.springframework.util.CollectionUtils;

import java.util.List;
import java.util.UUID;

/**
 * 参考Mybatis-Plus插件中的TenantSqlParser进行用户解析处理，
 * 其实现为使用jsqlparser对sql进行解析，拼装SQL语句
 * https://github.com/JSQLParser/JSqlParser
 * <p>
 * 用户数据隔离
 *
 * UUID问题难以解决: HexValue
 *
 * 会有一些莫名其妙的问题:
 * Caused by: net.sf.jsqlparser.parser.ParseException: Encountered unexpected token: "." "."
 *     at line 1, column 222.
 *
 * Was expecting one of:
 *
 *     "&"
 *     ","
 *     "::"
 *     ";"
 *     "<<"
 *     ">>"
 *     "ACTION"
 *     "ANY"
 *     "AS"
 *     "BYTE"
 *     "CASCADE"
 *     "CAST"
 *     "CHANGE"
 *     "CHAR"
 *     "CHARACTER"
 *     "COLLATE"
 *     "COLUMN"
 *     "COLUMNS"
 *     "COMMENT"
 *     "COMMIT"
 *     "CONNECT"
 *     "DESCRIBE"
 *     "DO"
 *     "DUPLICATE"
 *     "ENABLE"
 *     "END"
 *     "EXCEPT"
 *     "EXTRACT"
 *     "FALSE"
 *     "FIRST"
 *     "FN"
 *     "FOLLOWING"
 *     "FOR"
 *     "FROM"
 *     "GROUP"
 *     "HAVING"
 *     "INDEX"
 *     "INSERT"
 *     "INTERSECT"
 *     "INTERVAL"
 *     "INTO"
 *     "ISNULL"
 *     "KEY"
 *     "LAST"
 *     "MATERIALIZED"
 *     "MINUS"
 *     "NEXTVAL"
 *     "NO"
 *     "NULLS"
 *     "OPEN"
 *     "ORDER"
 *     "OVER"
 *     "PARTITION"
 *     "PATH"
 *     "PERCENT"
 *     "PRECISION"
 *     "PRIMARY"
 *     "PRIOR"
 *     "RANGE"
 *     "READ"
 *     "REPLACE"
 *     "ROW"
 *     "ROWS"
 *     "SEPARATOR"
 *     "SESSION"
 *     "SIBLINGS"
 *     "SIZE"
 *     "START"
 *     "TABLE"
 *     "TEMP"
 *     "TEMPORARY"
 *     "TOP"
 *     "TRUE"
 *     "TRUNCATE"
 *     "TYPE"
 *     "UNION"
 *     "UNSIGNED"
 *     "VALUE"
 *     "VALUES"
 *     "WHERE"
 *     "WINDOW"
 *     "XML"
 *     "ZONE"
 *     "["
 *     "^"
 *     "|"
 *     <EOF>
 *     <K_DATETIMELITERAL>
 *     <K_DATE_LITERAL>
 *     <S_CHAR_LITERAL>
 *     <S_IDENTIFIER>
 *     <S_QUOTED_IDENTIFIER>
 *
 * 	at net.sf.jsqlparser.parser.CCJSqlParser.generateParseException(CCJSqlParser.java:22439) ~[jsqlparser-3.1.jar:na]
 * 	at net.sf.jsqlparser.parser.CCJSqlParser.jj_consume_token(CCJSqlParser.java:22286) ~[jsqlparser-3.1.jar:na]
 * 	at net.sf.jsqlparser.parser.CCJSqlParser.Statements(CCJSqlParser.java:577) ~[jsqlparser-3.1.jar:na]
 * 	at net.sf.jsqlparser.parser.CCJSqlParserUtil.parseStatements(CCJSqlParserUtil.java:135) ~[jsqlparser-3.1.jar:na]
 * 	... 133 common frames omitted
 *
 * @author zhangfb
 * @since 2019/12/5
 */
@Slf4j
@Data
@Accessors(chain = true)
public class DataQueryInterceptor implements StatementInspector {

    private static final long serialVersionUID = -5090194680046440809L;

    /**
     * 当前用户ID，从UserContext获取
     */
    private static final ThreadLocal<UUID> USER_ID_CONTEXT = new ThreadLocal<>();

    /**
     * 需进行用户解析的用户字段名,本项目中为固定名称
     */
    private String userIdColumn = "created_by";

    /**
     * 重写StatementInspector的inspect接口，参数为hibernate处理后的原始SQL，
     * 返回值为我们修改后的SQL
     */
    @Override
    public String inspect(String sql) {
        String resSql = sql;
        // 用户信息为空返回原sql
        final UserPrincipal currentUser = CurrentUserContextHolder.getCurrentUser();
        if (currentUser == null) {
            return resSql;
        }
        try {
            // 设置用户
            USER_ID_CONTEXT.set(currentUser.getId());
            log.info("用户解析开始，原始SQL：{}", resSql);
            Statements statements = CCJSqlParserUtil.parseStatements(resSql);
            StringBuilder sqlStringBuilder = new StringBuilder();
            int i = 0;
            for (Statement statement : statements.getStatements()) {
                // 只处理查询请求
                if (statement instanceof Select) {
                    if (i++ > 0) {
                        sqlStringBuilder.append(';');
                    }
                    sqlStringBuilder.append(this.processParser(statement));
                }
            }
            resSql = sqlStringBuilder.toString();
            log.info("用户解析结束，解析后SQL：{}", resSql);
        } catch (Exception e) {
            log.error("用户解析失败，解析SQL异常", e);
            return resSql;
        } finally {
            USER_ID_CONTEXT.remove();
        }
        return resSql;
    }

    private String processParser(Statement statement) {
        if (statement instanceof Insert) {
            this.processInsert((Insert) statement);
        } else if (statement instanceof Select) {
            this.processSelectBody(((Select) statement).getSelectBody());
        } else if (statement instanceof Update) {
            this.processUpdate((Update) statement);
        } else if (statement instanceof Delete) {
            this.processDelete((Delete) statement);
        }
        //返回处理后的SQL
        return statement.toString();
    }

    /**
     * select 语句处理
     */

    private void processSelectBody(SelectBody selectBody) {
        if (selectBody instanceof PlainSelect) {
            processPlainSelect((PlainSelect) selectBody);
        } else if (selectBody instanceof WithItem) {
            WithItem withItem = (WithItem) selectBody;
            if (withItem.getSelectBody() != null) {
                processSelectBody(withItem.getSelectBody());
            }
        } else {
            SetOperationList operationList = (SetOperationList) selectBody;
            if (!CollectionUtils.isEmpty(operationList.getSelects())) {
                operationList.getSelects().forEach(this::processSelectBody);
            }
        }
    }

    /**
     * insert 语句处理
     */

    private void processInsert(Insert insert) {
        final Expression exception = this.uuidToHexValue(USER_ID_CONTEXT.get());
        insert.getColumns().add(new Column(userIdColumn));
        if (insert.getSelect() != null) {
            processPlainSelect((PlainSelect) insert.getSelect().getSelectBody(), true);
        } else if (insert.getItemsList() != null) {
            // fixed github pull/295
            ItemsList itemsList = insert.getItemsList();
            if (itemsList instanceof MultiExpressionList) {
                ((MultiExpressionList) itemsList).getExprList().forEach(el -> el.getExpressions().add(exception));
            } else {
                ((ExpressionList) insert.getItemsList()).getExpressions().add(exception);
            }
        } else {
            throw new AppRuntimeException("Failed to process multiple-table update, please exclude the tableName or statementId");
        }
    }

    /**
     * update 语句处理
     */

    private void processUpdate(Update update) {
        final Table table = update.getTable();
        update.setWhere(this.andExpression(table, update.getWhere()));
    }

    /**
     * delete 语句处理
     */

    private void processDelete(Delete delete) {
        delete.setWhere(this.andExpression(delete.getTable(), delete.getWhere()));
    }

    /**
     * delete update 语句 where 处理
     */
    private BinaryExpression andExpression(Table table, Expression where) {
        //获得where条件表达式
        EqualsTo equalsTo = new EqualsTo();
        equalsTo.setLeftExpression(this.getAliasColumn(table));
        equalsTo.setRightExpression(this.uuidToHexValue(USER_ID_CONTEXT.get()));
        if (null != where) {
            if (where instanceof OrExpression) {
                return new AndExpression(equalsTo, new Parenthesis(where));
            } else {
                return new AndExpression(equalsTo, where);
            }
        }
        return equalsTo;
    }

    /**
     * 处理 PlainSelect
     */
    private void processPlainSelect(PlainSelect plainSelect) {
        processPlainSelect(plainSelect, false);
    }

    /**
     * 处理 PlainSelect
     *
     * @param plainSelect ignore
     * @param addColumn   是否添加用户列,insert into select语句中需要
     */
    private void processPlainSelect(PlainSelect plainSelect, boolean addColumn) {
        FromItem fromItem = plainSelect.getFromItem();
        if (fromItem instanceof Table) {
            Table fromTable = (Table) fromItem;
            //#1186 github
            plainSelect.setWhere(builderExpression(plainSelect.getWhere(), fromTable));
            if (addColumn) {
                plainSelect.getSelectItems().add(new SelectExpressionItem(
                        new Column(userIdColumn)));
            }
        } else {
            processFromItem(fromItem);
        }
        List<Join> joins = plainSelect.getJoins();
        if (!CollectionUtils.isEmpty(joins)) {
            joins.forEach(j -> {
                processJoin(j);
                processFromItem(j.getRightItem());
            });
        }
    }

    /**
     * 处理子查询等
     */
    private void processFromItem(FromItem fromItem) {
        if (fromItem instanceof SubJoin) {
            SubJoin subJoin = (SubJoin) fromItem;
            if (subJoin.getJoinList() != null) {
                subJoin.getJoinList().forEach(this::processJoin);
            }
            if (subJoin.getLeft() != null) {
                processFromItem(subJoin.getLeft());
            }
        } else if (fromItem instanceof SubSelect) {
            SubSelect subSelect = (SubSelect) fromItem;
            if (subSelect.getSelectBody() != null) {
                processSelectBody(subSelect.getSelectBody());
            }
        } else if (fromItem instanceof ValuesList) {
            log.debug("Perform a subquery, if you do not give us feedback");
        } else if (fromItem instanceof LateralSubSelect) {
            LateralSubSelect lateralSubSelect = (LateralSubSelect) fromItem;
            if (lateralSubSelect.getSubSelect() != null) {
                SubSelect subSelect = lateralSubSelect.getSubSelect();
                if (subSelect.getSelectBody() != null) {
                    processSelectBody(subSelect.getSelectBody());
                }
            }
        }
    }

    /**
     * 处理联接语句
     */
    private void processJoin(Join join) {
        if (join.getRightItem() instanceof Table) {
            Table fromTable = (Table) join.getRightItem();
            join.setOnExpression(builderExpression(join.getOnExpression(), fromTable));
        }
    }

    /**
     * 处理条件:
     * 支持 getTenantHandler().getUserId()是一个完整的表达式：tenant in (1,2)
     * 默认tenantId的表达式： LongValue(1)这种依旧支持
     */
    private Expression builderExpression(Expression currentExpression, Table table) {
        final Expression expression = this.uuidToHexValue(USER_ID_CONTEXT.get());
        EqualsTo appendExpression = new EqualsTo();
        appendExpression.setLeftExpression(this.getAliasColumn(table));
        appendExpression.setRightExpression(expression);
        if (currentExpression == null) {
            return appendExpression;
        }
        if (currentExpression instanceof BinaryExpression) {
            BinaryExpression binaryExpression = (BinaryExpression) currentExpression;
            doExpression(binaryExpression.getLeftExpression());
            doExpression(binaryExpression.getRightExpression());
        } else if (currentExpression instanceof InExpression) {
            InExpression inExp = (InExpression) currentExpression;
            ItemsList rightItems = inExp.getRightItemsList();
            if (rightItems instanceof SubSelect) {
                processSelectBody(((SubSelect) rightItems).getSelectBody());
            }
        }
        if (currentExpression instanceof OrExpression) {
            return new AndExpression(new Parenthesis(currentExpression), appendExpression);
        } else {
            return new AndExpression(currentExpression, appendExpression);
        }
    }

    private void doExpression(Expression expression) {
        if (expression instanceof FromItem) {
            processFromItem((FromItem) expression);
        } else if (expression instanceof InExpression) {
            InExpression inExp = (InExpression) expression;
            ItemsList rightItems = inExp.getRightItemsList();
            if (rightItems instanceof SubSelect) {
                processSelectBody(((SubSelect) rightItems).getSelectBody());
            }
        }
    }

    /**
     * 目前: 针对自定义的tenantId的条件表达式[tenant_id in (1,2,3)]，无法处理多用户的字段加上表别名
     * select a.id, b.name
     * from a
     * join b on b.aid = a.id and [b.]tenant_id in (1,2) --别名[b.]无法加上 TODO
     *
     * @param expression
     * @param table
     * @return 加上别名的多用户字段表达式
     */
    private Expression processTableAlias4CustomizedTenantIdExpression(Expression expression, Table table) {
        //cannot add table alias for customized userId expression,
        // when tables including userId at the join table poistion
        return expression;
    }

    /**
     * 用户字段别名设置
     * <p>tableName.userId 或 tableAlias.userId</p>
     *
     * @param table 表对象
     * @return 字段
     */
    private Column getAliasColumn(Table table) {
        StringBuilder column = new StringBuilder();
        if (null == table.getAlias()) {
            column.append(table.getName());
        } else {
            column.append(table.getAlias().getName());
        }
        column.append(".");
        column.append(userIdColumn);
        return new Column(column.toString());
    }

    private HexValue uuidToHexValue(UUID userId) {
        return new UUIDValue(userId);
    }
}